library(reshape)
library(gridExtra)
library(ggplot2)
require(reshape2)
library(tidyverse)
library(xtable)
rm(list=ls())

## For publication 
grayscale <- TRUE

load('simresults_newbody.Rda')

colnames(output$point)[1]<-colnames(output$ses)[1]<-"n_size"
colnames(output$point)[8]<-colnames(output$ses)[8]<-"PLCE"
colnames(output$point)[9]<-colnames(output$ses)[9]<-"GRF"

options(device="quartz")


i.mean<-i.err<-i.inter<-F

gen.cover<-function(n.use,i.mean,i.err,i.inter,i.RE,col.use){
theta.run<-output$point
se.run<-output$ses
sub.curr<-
  theta.run[,"n_size"]==n.use  &
  theta.run[,"meanhet"]==i.mean&theta.run[,"errhet"]==i.err&theta.run[,"inter"]==i.inter&
  theta.run[,"REs"]==i.RE
theta.sub<-output$point[sub.curr,]
sub.curr<-
  se.run[,"n_size"]==n.use  &
  se.run[,"meanhet"]==i.mean&se.run[,"errhet"]==i.err&se.run[,"inter"]==i.inter
se.sub<-output$ses[sub.curr,]
#se.sub[,"PLCE"]<-se.sub[,"PLCE"]
alpha<-seq(0.01,0.99,by=0.01)
sapply(alpha,FUN=function(z){
cv<--qnorm((1-z)/2)
out1<-mean(abs(theta.sub[,col.use]-1)<cv*se.sub[,col.use])
if(out1==0) out1<-abs(rnorm(1,sd=0.01))
out1
})
}

## Emulate ggplot2 colors ----
gg_color <- function(n,alpha0=1) {
  if(grayscale) {
   # gc<- gray.colors(n, start=0, end =1, alpha=alpha0)[1:3]
   gc<- gray.colors(4, start=0, end =1, alpha=alpha0)[1:3]
   return(c(gc,rev(gc)))
  }
  hues = seq(15, 375, length = n + 1)
  hcl(h = hues, l = 65, c = 100,alpha=alpha0)[1:n]
  
}

## Function for coverage plot ----
make.coverplot<-function(n.use,i.mean,i.err,i.inter,i.RE){
  #"DML"     "HOE"     "RF"      "CBPS"    "OLS"     "KRLS" 
  truecover<-seq(0.01,0.99,by=0.01)
  names.vec<-c("PLCE","GRF","DML","CBPS","KRLS","OLS")
  cover.mat<-matrix(NA,nrow=length(truecover),ncol=6)
  colnames(cover.mat)<-names.vec
  for(i in 1:6){
    cover.mat[,i]<-gen.cover(n.use, i.mean, i.err, i.inter, i.RE, names.vec[i])
  }

  plot(0,0,xlim=c(0,1),ylim=c(0,1),type="n",xaxt="n",yaxt="n",xlab="",ylab="",bty="n")
  axis(1,at=(0:10)/10,col = "white", cex.axis=1.25 )
  axis(2,at=(0:10)/10,col = "white", cex.axis=1.25)
  mtext(side=1,"Expected Coverage",font=2,line=1.9, cex=1.2)
  mtext(side=2,"Actual Coverage",font=2,line=2.2, las=0, cex=1.2)
  
  if(grayscale)
    rect(par("usr")[1],par("usr")[3],par("usr")[2],par("usr")[4],col = "gray95")
    else
    rect(par("usr")[1],par("usr")[3],par("usr")[2],par("usr")[4],col = "gray90")
  abline(v=(1:9)/10,col="white")
  abline(h=(1:9)/10,col="white")
  
  abline(c(0,1),col=gray(0.1),lwd=1,lty=1)
  for(i in 1:6){
  lines(truecover,cover.mat[,i],lwd=2,lty=1+grayscale*(i>3),col=gg_color(6)[i])
  }
  legend("topleft",legend=colnames(cover.mat),col=gg_color(6),lty=c(1,1,1,2,2,2),xpd=T,lwd=3,bg="gray90",bty="n",box.col="gray90", seg.len = 4, cex=1.25)
  
}


## Function for box and whisker plot ----
make.bwplot<-function(n.use,i.mean,i.err,i.inter,i.REs){
  #"DML"     "PLCE"     "RF"      "CBPS"    "OLS"     "KRLS" 

  output2<-data.frame(1,output$point)
  output2<-tibble(output2)
  sub.use<-output2%>% filter(n_size==n.use,meanhet==i.mean,errhet==i.err,inter==i.inter,REs==i.REs)%>%dplyr::select(PLCE,GRF,DML,CBPS,KRLS,OLS) 
  ylim.use<-range(output2%>%filter(inter==i.inter)%>% dplyr::select(PLCE,GRF,DML,CBPS,KRLS,OLS))
  if(i.inter==0){
  ylim.use[2]<-3.2
  ylim.use[1]<-.15
  } else {
    ylim.use[2]<-3.2
    ylim.use[1]<-.15
  }
  plot(0,0,xlim=c(0,7),ylim=ylim.use,type="n",xaxt="n",yaxt="n",xlab="",ylab="",bty="n")
  #axis(1,at=(0:10),col = "white")
  axis(2,at=(0:10),col = "white")
  mtext(side=2,"Estimate",font=2,line=2, las=0)
  if(grayscale)
    rect(par("usr")[1],par("usr")[3],par("usr")[2],par("usr")[4],col = "gray95")
  else
    rect(par("usr")[1],par("usr")[3],par("usr")[2],par("usr")[4],col = "gray90")
  
  abline(v=(1:9),col="white")
  abline(h=(1:9),col="white")
  abline(h=1,lwd=3,col="gray40")
  
  if(grayscale)
    boxplot(sub.use,col = gray(.8, alpha=.8),add=TRUE,cex.axis=1.25,pch="",ylab="",yaxt="n")
  else
    boxplot(sub.use,col = gg_color(6,alpha0=.8),add=TRUE,cex=.5,pch="",ylab="",yaxt="n")
  

}


## PDF for figure in body with interference ----
pdf(h=12*.75,w=12*.75,"Fig2_Ratkovic.pdf")

par(mar = c(3, 3, 2, 2), # Dist' from plot to side of page
    mgp = c(2, 0.4, 0), # Dist' plot to label
    las = 1, # Rotate y-axis text
    tck = -.01, # Reduce tick length
    xaxs = "i", yaxs = "i", # Remove plot padding
    oma =c(0,5.2,0,0),#c(2,2,2,2),
    mfrow=c(4,2))

  for(i.meanhet in c(F,T)){
    for(i.REs in c(F,T)) {
      if(i.meanhet){ i.mean <- i.err <- T} else {i.mean<-i.err <-F}
  make.bwplot(1000,i.mean,i.err,i.REs=i.REs,i.inter=TRUE)
  make.coverplot(1000,i.mean, i.err, i.RE=i.REs,i.inter=TRUE)
  }}

mtext(outer=T,c("\nBoth", "Treatment Effect\n Heterogeneity","\nRandom Effects ",
                "\nBaseline" ),at=c(0.125,.38,.635,.89),side=2,las=0,line=1.3,cex=1.4,font=2)
dev.off()


## PDF for figure in body without interference ----
pdf(h=12*.75,w=12*.75,"Fig1_Ratkovic.pdf")

par(mar = c(3, 3, 2, 2), # Dist' from plot to side of page
    mgp = c(2, 0.4, 0), # Dist' plot to label
    las = 1, # Rotate y-axis text
    tck = -.01, # Reduce tick length
    xaxs = "i", yaxs = "i", # Remove plot padding
    oma =c(0,5.2,0,0),#c(2,2,2,2),
    mfrow=c(4,2))

for(i.meanhet in c(F,T)){
  for(i.REs in c(F,T)) {
    if(i.meanhet){ i.mean <- i.err <- T} else {i.mean<-i.err <-F}
    make.bwplot(1000,i.mean,i.err,i.REs=i.REs,i.inter=FALSE)
    make.coverplot(1000,i.mean, i.err, i.RE=i.REs,i.inter=FALSE)
    if(i.REs==F&i.meanhet==F) points(0.9,0.845,pch=4,cex=1.4)
    
  }}

mtext(outer=T,c("\nBoth", "Treatment Effect\n Heterogeneity","\nRandom Effects ",
                "\nBaseline" ),at=c(0.125,.38,.635,.89),side=2,las=0,line=1.3,cex=1.4,font=2)
dev.off()




